package rddl.policy.expectation;

import org.apache.commons.math3.random.RandomDataGenerator;
import rddl.EvalException;
import rddl.RDDL;
import rddl.State;
import util.Pair;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Random;

import rddl.policy.expectation.Util;

/**
 * 表示使用最大期望的State，在MaxExpectationPolicy中使用
 * <p>
 * Created by HuGuodong on 2017/11/24.
 */

public class MEPState extends State {

    public final static boolean DISPLAY_UPDATES = false;

    /**
     * for now:
     * compute the following state: (s'|s,a)
     *
     * @param actions
     * @param _rand
     * @throws EvalException
     */
    public double computeNextStateReward(ArrayList<RDDL.PVAR_INST_DEF> actions, RandomDataGenerator _rand) throws EvalException {

        // modify RDDL.Bernoulli sample method to return double or boolean value.
        RDDL.SWITCH_TO_SAMPLE = false;

        // Clear then set the actions
        for (RDDL.PVAR_NAME p : _actions.keySet())
            _actions.get(p).clear();
        setPVariables(_actions, actions);

        //System.out.println("Starting state: " + _state + "\n");
        //System.out.println("Starting nonfluents: " + _nonfluents + "\n");

        // First compute intermediate variables (derived should have already been computed)
        HashMap<RDDL.LVAR, RDDL.LCONST> subs = new HashMap<RDDL.LVAR, RDDL.LCONST>();

        if (DISPLAY_UPDATES) System.out.println("Updating next state");
        for (RDDL.PVAR_NAME p : _alStateNames) {

            // Get default value
            Object def_value = null;
            RDDL.PVARIABLE_DEF pvar_def = _hmPVariables.get(p);
            if (!(pvar_def instanceof RDDL.PVARIABLE_STATE_DEF) ||
                    ((RDDL.PVARIABLE_STATE_DEF) pvar_def)._bNonFluent)
                throw new EvalException("Expected state variable, got nonfluent: " + p);
            def_value = ((RDDL.PVARIABLE_STATE_DEF) pvar_def)._oDefValue;

            // Generate updates for each ground fluent
            RDDL.PVAR_NAME primed = new RDDL.PVAR_NAME(p._sPVarName + "'");
            //System.out.println("Updating next state var " + primed + " (" + p + ")");
            ArrayList<ArrayList<RDDL.LCONST>> gfluents = generateAtoms(p);

            for (ArrayList<RDDL.LCONST> gfluent : gfluents) {
                if (DISPLAY_UPDATES) System.out.print("- " + primed + gfluent + " := ");
                RDDL.CPF_DEF cpf = _hmCPFs.get(primed);
                if (cpf == null)
                    throw new EvalException("Could not find cpf for: " + primed +
                            "... did you forget to prime (') the variable in the cpf definition?");

                subs.clear();
                for (int i = 0; i < cpf._exprVarName._alTerms.size(); i++) {
                    RDDL.LVAR v = (RDDL.LVAR) cpf._exprVarName._alTerms.get(i);
                    RDDL.LCONST c = (RDDL.LCONST) gfluent.get(i);
                    subs.put(v, c);
                }

                // 每个state(?x)拿到cpf的表达式: cpf $c1, $c2，running'($c1), running'($c2)可能是不同的
                // FIXME cpf._exprEqual 这个表达式需要修改，改为返回值为double类型
                Object value = cpf._exprEquals.sample(subs, this, _rand); // ?x:$c1 ...

                // FIXME 模拟返回double类型值
                if (value instanceof Boolean) {
                    // boolean -> 1.0/0.0, in fact,
                    // it's trivial, because boolean will be transform into double value, when do calculation.
                    value = ((Boolean) value) ? 1.0 : 0.0;
                }

                // monitor value
                System.out.println(cpf._exprVarName._pName +
                        "(" + subs.values().toString().replaceAll("\\[|\\]", "") + ")" + " = " + value);

                // Update value if not default
                // comment: 更新下一个状态值
                if (!value.equals(def_value)) {
                    HashMap<ArrayList<RDDL.LCONST>, Object> pred_assign = _nextState.get(p);
                    pred_assign.put(gfluent, value);
                }
            }
        }

        // 计算下一个状态 reward
        this._state = this._nextState;
        Double rewardExceptation = RDDL.ConvertToNumber(
                this._reward.sample(new HashMap<RDDL.LVAR, RDDL.LCONST>(), this, _rand)).doubleValue();

        System.out.println("MEPState next state's reward expectation: " + rewardExceptation);

        // modify RDDL.Bernoulli sample method to return double or boolean value.
        RDDL.SWITCH_TO_SAMPLE = true;

        return rewardExceptation;

    }


}
